from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import os
import logging
from dotenv import load_dotenv

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 加载环境变量
load_dotenv()

# 数据库连接配置
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./pilot_manage.db")
logger.info(f"数据库连接URL: {DATABASE_URL.split('://')[0]}://****:****@{DATABASE_URL.split('@')[1] if '@' in DATABASE_URL else 'localhost'}")

# 创建数据库引擎
try:
    # 如果是MySQL数据库，添加一些特定配置
    if DATABASE_URL.startswith('mysql'):
        logger.info("正在连接MySQL数据库...")
        # 尝试先创建数据库（如果不存在）
        try:
            import pymysql
            # 解析数据库连接参数
            db_params = {
                'host': DATABASE_URL.split('@')[1].split(':')[0],
                'port': int(DATABASE_URL.split('@')[1].split(':')[1].split('/')[0]),
                'user': DATABASE_URL.split('://')[1].split(':')[0],
                'password': DATABASE_URL.split('://')[1].split(':')[1].split('@')[0],
                'database': DATABASE_URL.split('/')[-1]
            }
            
            logger.info(f"尝试连接到MySQL服务器: {db_params['host']}:{db_params['port']}")
            # 创建数据库连接（不带数据库名）
            conn = pymysql.connect(
                host=db_params['host'],
                port=db_params['port'],
                user=db_params['user'],
                password=db_params['password'],
                charset='utf8mb4'
            )
            cursor = conn.cursor()
            # 创建数据库（如果不存在）
            cursor.execute(f"CREATE DATABASE IF NOT EXISTS {db_params['database']}")
            cursor.close()
            conn.close()
            logger.info(f"数据库 {db_params['database']} 检查/创建成功")
        except Exception as e:
            logger.warning(f"创建数据库失败: {e}，将尝试直接连接")
        
        # 创建数据库引擎
        engine = create_engine(
            DATABASE_URL,
            pool_size=10,
            max_overflow=20,
            pool_recycle=3600,
            pool_pre_ping=True,
            pool_timeout=30,
            echo=False,
            connect_args={
                'charset': 'utf8mb4',
                'connect_timeout': 10
            }
        )
        logger.info("MySQL数据库引擎创建成功")
    else:
        # SQLite配置
        logger.info("使用SQLite数据库")
        engine = create_engine(
            DATABASE_URL,
            pool_size=10,
            max_overflow=20,
            pool_recycle=3600,
            pool_pre_ping=True
        )
        logger.info("SQLite数据库引擎创建成功")
except Exception as e:
    logger.error(f"创建数据库引擎失败: {e}")
    raise

# 创建数据库会话
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# 创建基础模型类
Base = declarative_base()

# 获取数据库会话的依赖项
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()